"""
UNet训练配置
"""

# 数据配置
DATA_CONFIG = {
    'images_dir': 'data/images',
    'masks_dir': 'data/masks',
    'image_size': 256,
    'to_rgb': False,  # UNet可以直接处理单通道，不需要转3通道
}

# 训练配置
TRAIN_CONFIG = {
    'batch_size': 24,
    'num_epochs': 50,
    'learning_rate': 1.5e-4,
    'weight_decay': 1e-5,
    'num_workers': 8,
    'k_folds': 8,
    'seed': 42,
}

# 模型配置
MODEL_CONFIG = {
    'in_channels': 1,
    'out_channels': 1,
    'features': [64, 128, 256, 512],
}

# 损失函数配置
LOSS_CONFIG = {
    'bce_weight': 1.0,
    'dice_weight': 1.0,
}

# 优化器配置
OPTIMIZER_CONFIG = {
    'type': 'adam',  # 'adam' or 'sgd'
}

# 学习率调度器配置
SCHEDULER_CONFIG = {
    'type': 'cosine',  # 'step' or 'cosine'
    'step_size': 30,
    'gamma': 0.1,
    'T_max': 50,
    'eta_min': 1e-6,
}

# 其他配置
OTHER_CONFIG = {
    'device': 'cuda' if __import__('torch').cuda.is_available() else 'cpu',
    'save_interval': 10,  # 每隔多少epoch保存一次
    'log_interval': 10,  # 每隔多少iter打印一次
}
